import numpy as np
import os
import pandas as pd
from sklearn.metrics import precision_recall_fscore_support
from strategy.datas.data import XDataset, YDataset
from strategy.rl.rl_model1 import RLFitModel
from tqdm import tqdm


def main_job():
    savedir = "rub9"
    os.makedirs(savedir, exist_ok=True)
    delta = 10
    # df = pd.read_parquet("cache/1s/BTC_USDT/1704038400000_1731081600000.parquet")
    df = pd.read_parquet("../cache/3m/ETH_USDT/1609430400000_1731081600000.parquet")
    df = df.iloc[:2000]

    split_idx = int(len(df) * 0.8)
    train_idx = df.index[100:split_idx]
    test_idx = df.index[split_idx + 200 : -50]

    train_df = df.iloc[: split_idx + 50].copy(deep=True)
    test_df = df.iloc[split_idx:].copy(deep=True)

    train_set = XDataset(train_df, train_idx)
    test_set = XDataset(test_df, test_idx)

    train_y = YDataset(train_df, train_idx)
    test_y = YDataset(test_df, test_idx)

    X = np.array(train_set)
    y = np.array(train_y)
    unique_values, counts = np.unique(y, return_counts=True)
    print(f"数值和数量 {unique_values} {counts}")
    print(f"训练标签标均值准差: {y.mean()} {y.std()}")

    t_X = np.array(test_set)
    t_y = np.array(test_y)

    N = len(train_y)
    is_classify = False

    if is_classify:
        weights = {k: N / (v + 1) for k, v in zip(unique_values, counts)}
        # print(f"fit {weights}")

    # model = XGBRegressor(
    #     n_estimators=20,  # 树的数量
    #     learning_rate=0.1,  # 学习率
    #     max_depth=5,       # 树的最大深度
    #     random_state=42,    # 随机种子
    # )
    # model = XGBClassifier(
    #     n_estimators=200,  # 树的数量
    #     learning_rate=0.1,  # 学习率
    #     max_depth=5,  # 树的最大深度
    #     random_state=42,  # 随机种子
    # )
    # model = GradientBoostingClassifier(n_estimators=20, max_depth=5, verbose=10)
    print("start fitting")
    # model = CatBoostRegressor(
    #     iterations=1000,  # 最大迭代次数
    #     learning_rate=0.03,  # 学习率
    #     depth=4,  # 树的深度
    #     # n_estimators=200,
    #     # cat_features=cat_features,  # 指定类别特征
    #     verbose=100,  # 每隔 100 次输出一次日志
    #     early_stopping_rounds=50,  # 早停防止过拟合
    #     l2_leaf_reg=10,
    #     bagging_temperature=1.0,
    #     subsample=0.8,
    #     # class_weights=weights
    # )

    # model = CatBoostClassifier(
    #     iterations=500,  # 最大迭代次数
    #     learning_rate=0.03,  # 学习率
    #     depth=5,  # 树的深度
    #     # n_estimators=200,
    #     # cat_features=cat_features,  # 指定类别特征
    #     verbose=100,  # 每隔 100 次输出一次日志
    #     early_stopping_rounds=50,  # 早停防止过拟合
    #     l2_leaf_reg=10,
    #     bagging_temperature=1.0,
    #     # subsample=0.8
    #     class_weights=weights
    # )

    model = RLFitModel()
    # model = RollingFit(TrainModel2())
    # model = TrainModel()
    # model.fit(X, y)
    # model = RollingFit(model)

    y_cls = y
    model.fit(X, y_cls)  # , epochs=5, lr = 1e-4
    model.save_model(os.path.join(savedir, "model.data"))
    # model.load_model(os.path.join(savedir, "model.data"))
    results = model.predict(X).squeeze()
    if not is_classify:
        y_cls = np.sign(y_cls) + 1
        results = np.sign(results) + 1

    precision, recall, f1, support = precision_recall_fscore_support(
        y_cls[results != 1], results[results != 1], average=None
    )
    print(f"训练集 precision recall f1 support {precision} {recall} {f1} {support}")

    test_results = model.predict(t_X).squeeze()
    if not is_classify:
        t_y = np.sign(t_y) + 1
        test_results = np.sign(test_results) + 1

    precision, recall, f1, support = precision_recall_fscore_support(
        t_y[test_results != 1], test_results[test_results != 1], average=None
    )
    print(f"测试集 precision recall f1 support {precision} {recall} {f1} {support}")
    print(
        f"预测结果方差 {np.std(test_results)}, 训练集方差 {np.std(y)}, 训练集预测方差：{np.std(results)}"
    )

    # with open(os.path.join(savedir, "weights.json"), "w") as f:
    #     json.dump(weights, f)
    test_results = np.where(test_results == 2, -1, test_results)

    input("wait")

    pre = 0
    hold = 0
    s = 0
    trades = 0
    slip = 0.0000
    last_price = None
    last_achieve = None
    poscnt = 0

    K = 500000

    test_df.loc[test_idx, "result"] = test_results
    for i, rowidx in tqdm(enumerate(test_idx)):
        if i > K:
            K = i
            break

        hold += 1
        profit_rate = None
        cur_price = test_df.loc[rowidx, "close"]
        cur_result = test_df.loc[rowidx, "result"]
        if last_price:
            profit_rate = (cur_price - last_price) / last_price * pre

        if np.isfinite(cur_result) and cur_result > 0.00005:
            confirm = 1
        elif np.isfinite(cur_result) and cur_result < -0.00005:
            confirm = -1
        else:
            confirm = 0

        if confirm == 1:

            if pre == 0:
                test_df.loc[rowidx, "op"] = 1
                pre = 1
                s -= cur_price * (1 + slip)
                trades += 1
                if profit_rate is not None and profit_rate > 0:
                    poscnt += 1
                hold = 0
                last_price = cur_price

                profit_rate = None
                last_achieve = None
            elif pre == -1:
                test_df.loc[rowidx, "op"] = 1
                s -= 2 * cur_price * (1 + slip)
                pre = 1
                trades += 1
                if profit_rate is not None and profit_rate > 0:
                    poscnt += 1
                hold = 0
                last_price = cur_price
                profit_rate = None
                last_achieve = None
        elif confirm == -1:
            if pre == 0:
                test_df.loc[rowidx, "op"] = -1
                pre = -1
                s += cur_price * (1 - slip)
                trades += 1
                if profit_rate is not None and profit_rate > 0:
                    poscnt += 1
                hold = 0
                last_price = cur_price
                profit_rate = None
                last_achieve = None

            elif pre == 1:
                test_df.loc[rowidx, "op"] = -1
                s += 2 * cur_price * (1 - slip)
                pre = -1
                trades += 1
                if profit_rate is not None and profit_rate > 0:
                    poscnt += 1
                hold = 0
                last_price = cur_price
                profit_rate = None
                last_achieve = None
        if (
            confirm == 2
            or hold >= 20 * delta
            or (
                profit_rate is not None
                and (
                    False
                    # profit_rate > 2e-2
                    # profit_rate < -5e-3 or
                    # (last_achieve is not None and profit_rate - last_achieve < -5e-3)
                )
            )
        ):  # 持仓时间达到要求
            # if (last_achieve and profit_rate - last_achieve < -10 * target_change):
            #     print("debug")

            if pre == 1:
                test_df.loc[rowidx, "op"] = -1
                s += cur_price * (1 - slip)
                pre = 0
                trades += 1
                if profit_rate is not None and profit_rate > 0:
                    poscnt += 1
                profit_rate = None
                last_achieve = None
            elif pre == -1:
                test_df.loc[rowidx, "op"] = 1
                s -= cur_price * (1 + slip)
                pre = 0
                trades += 1
                if profit_rate is not None and profit_rate > 0:
                    poscnt += 1
                profit_rate = None
                last_achieve = None
        if last_achieve is not None:
            last_achieve = max(last_achieve, profit_rate)
        else:
            last_achieve = profit_rate

    K = min(K, len(test_idx) - 1)
    if pre == 1:
        s += test_df.loc[test_idx[K], "close"] * (1 - slip)
    elif pre == -1:
        s -= test_df.loc[test_idx[K], "close"] * (1 + slip)
    print("总共收益,交易次数，胜率", s, trades, poscnt / trades)


if __name__ == "__main__":
    main_job()
